use std::{net::SocketAddr, sync::Arc};

use axum::{
    extract::{
        ws::{Message, WebSocket},
        ConnectInfo, State, TypedHeader, WebSocketUpgrade,
    },
    headers,
    response::IntoResponse,
    routing::get,
    Router,
};
use futures::{SinkExt, StreamExt};

use crate::{config::Config, state::create_state};

pub fn create_router(cfg: Config) -> Router {
    Router::new()
        .route("/ws", get(ws_handler))
        .with_state(create_state(cfg))
}

pub async fn ws_handler(
    State(state): State<Arc<crate::state::State>>,
    ws: WebSocketUpgrade,
    user_agent: Option<TypedHeader<headers::UserAgent>>,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
) -> impl IntoResponse {
    let user_agent = if let Some(TypedHeader(ua)) = user_agent {
        ua.to_string()
    } else {
        String::from("Unknown user agent")
    };

    ws.on_upgrade(move |socket| handle_socket(socket, addr, user_agent, state))
}

async fn handle_socket(
    socket: WebSocket,
    addr: SocketAddr,
    user_agent: String,
    state: Arc<crate::state::State>,
) {
    state
        .clone()
        .emitter
        .lock()
        .await
        .track(socket, addr, user_agent, state)
        .await;
}
